import numpy as np
from dataset import *

import cPickle
import logging
import os
import pdb

def load_roidb(roidb_path_list, imglst_path_list=None, filter_strategy=None,\
    load_proposals=False, proposal_file_list=None, percent=None, startInd=None):
    """
    load roidb from cache
    """
    roidb_list  = []
    if roidb_path_list is not None:
        for roidb_path in roidb_path_list:
            with open(roidb_path, 'rb') as fid:
                roidb   = cPickle.load(fid)
            roidb_list.append(roidb)
    assert len(roidb_list) == len(imglst_path_list)
    if percent is not None:
        assert 0 < percent and percent <= 1.0, "percent:{} must in (0,1]".format(percent)
        print("load roidb percent is:{}".format(percent))
        logging.info("load roidb percent is:{}".format(percent))
        new_roidb_list  = []
        for roidb in roidb_list:
            all_roidb_len   = len(roidb)
            use_roidb_len   = int(all_roidb_len * percent)
            if startInd is None:
                startInd    = 0
            new_roidb       = roidb[startInd:startInd+use_roidb_len]
            new_roidb_list.append(new_roidb)
        # replace the original roidb_list
        roidb_list  = new_roidb_list
    if load_proposals and len(proposal_file_list) > 0:
        # for sniper proposal using
        # TODO
        pass
    if filter_strategy is not None:
        roidb_list  = [filter_roidb_new(roidb, filter_strategy) for roidb in roidb_list]
    if imglst_path_list is not None:
        add_roidb_rec_index(roidb_list, imglst_path_list)
    roidb   = merge_roidb(roidb_list)
    logging.info('total num images: %d' % len(roidb))
    return roidb
        
def add_roidb_rec_index(roidb_list, imglst_path_list):
    assert len(roidb_list) == len(imglst_path_list)
    for i, roidb in enumerate(roidb_list):
        img_list = {}
        with open(imglst_path_list[i]) as fin:
            for line in iter(fin.readline, ''):
                line = line.strip().split('\t')
                img_list[line[-1]] = int(line[0])
        for roi_rec in roidb:
            img_name = roi_rec['image']
            if img_name not in img_list:
                img_name = os.path.basename(roi_rec['image'])
                assert img_name in img_list
            roi_rec['rec_index'] = img_list[img_name]
            roi_rec['imgrec_id'] = i

def filter_roidb_new(roidb, filter_strategy, need_inds=False):
    """
    from smart
    """
    all_choose_inds = range(len(roidb))

    def filter_roidb_func(choose_inds, filter_name, filter_func):
        if filter_name in filter_strategy and filter_strategy[filter_name]:
            num = len(choose_inds)
            choose_inds = [i for i in choose_inds if not filter_func(roidb[i])]
            num_after = len(choose_inds)
            logging.info('filter %d %s roidb entries: %d -> %d' % (num - num_after, filter_name[7:], num, num_after))
        return choose_inds

    def is_points_as_boxes(entry):
        gt_boxes = entry['boxes']
        width = gt_boxes[:, 2] - gt_boxes[:, 0] + 1
        height = gt_boxes[:, 3] - gt_boxes[:, 1] + 1
        flag = (width > 1).all() and (height > 1).all()
        return not flag
    all_choose_inds = filter_roidb_func(all_choose_inds, 'remove_point', is_points_as_boxes)

#    def is_valid(entry):
#        """
#        borrowed from filter_roidb
#        valid images have at least 1 fg or bg rois
#        """
#        overlaps= entry['max_overlaps']
#        fg_inds = np.where(overlaps >= config.TRAIN.FG_THRESH)[0]
#        bg_inds = np.where((overlaps < config.TRAIN.BG_THRESH_HI) & (overlaps > config.TRAIN.BG_THRESH_LO))[0]
#        # valid   = len(fg_inds) > 0 or len(bg_inds) > 0
#        valid   = len(fg_inds) > 0
#        return valid
#    all_choose_inds = filter_roidb_func(all_choose_inds, 'is_valid', is_valid)

    def is_empty_boxes(entry):
        num_valid_boxes = np.sum(entry['gt_classes'] > 0)
        return num_valid_boxes == 0
    all_choose_inds = filter_roidb_func(all_choose_inds, 'remove_empty_boxes', is_empty_boxes)

    def is_single_boxes(entry):
        num_valid_boxes = np.sum(entry['gt_classes'] > 0)
        return num_valid_boxes <= 1
    all_choose_inds = filter_roidb_func(all_choose_inds, 'remove_single_boxes', is_single_boxes)

    def is_multi_boxes(entry):
        num_valid_boxes = np.sum(entry['gt_classes'] > 0)
        return num_valid_boxes > 1
    all_choose_inds = filter_roidb_func(all_choose_inds, 'remove_multi_boxes', is_multi_boxes)

    def is_empty_kps(entry):
        scores = entry['keypoints'][:, 2::3]
        scores = np.sum(scores, axis=1)
        keep = np.where(scores != 0)[0]
        entry['keypoints'] = entry['keypoints'][keep, :]
        entry['boxes'] = entry['boxes'][keep, :]
        if 'gt_classes' in entry:
            entry['gt_classes'] = entry['gt_classes'][keep]
        return np.sum(scores) == 0
    all_choose_inds = filter_roidb_func(all_choose_inds, 'remove_empty_kps', is_empty_kps)

    def is_any_unvis_kps(entry):
        scores = entry['keypoints'][:, 2::3]
        return (scores != 2).any()
    all_choose_inds = filter_roidb_func(all_choose_inds, 'remove_any_unvis_kps', is_any_unvis_kps)

    if 'max_num_images' in filter_strategy:
        max_num_images = filter_strategy['max_num_images']
        if 0 < max_num_images < len(all_choose_inds):
            num = len(all_choose_inds)
            all_choose_inds = all_choose_inds[:max_num_images]
            num_after = len(all_choose_inds)
            logging.info('filter %d roidb entries after max_num_images: %d -> %d' % (num - num_after, num, num_after))

    if 'parts' in filter_strategy:
        part_index = filter_strategy['parts'][0]
        num_parts = filter_strategy['parts'][1]
        assert part_index < num_parts
        num_inds_per_part = (len(all_choose_inds) + num_parts - 1) / num_parts
        num = len(all_choose_inds)
        all_choose_inds = all_choose_inds[part_index*num_inds_per_part: (part_index+1)*num_inds_per_part]
        num_after = len(all_choose_inds)
        logging.info('filter %d roidb entries after parts: %d -> %d' % (num - num_after, num, num_after))

    if 'indexes' in filter_strategy:
        start_index = filter_strategy['indexes'][0]
        end_index = filter_strategy['indexes'][1]
        num = len(all_choose_inds)
        assert 0 <= start_index < end_index <= num
        all_choose_inds = all_choose_inds[start_index:end_index]
        num_after = len(all_choose_inds)
        logging.info('filter %d roidb entries after indexes: %d -> %d' % (num - num_after, num, num_after))

    roidb = [roidb[i] for i in all_choose_inds]

    if need_inds:
        return roidb, all_choose_inds
    else:
        return roidb

def load_merge_filter_gt_roidb_wrapper(config):
    """
    wrapper for load_gt_roidb, merge_roidb, filter_roidb using config
    """
    image_sets  = [iset for iset in config.dataset.image_set.split("+")]
    roidbs      = [load_gt_roidb(config.dataset.dataset, image_set, config.dataset.root_path,
            config.dataset.dataset_path, flip=config.TRAIN.FLIP) 
            for image_set in image_sets]
    roidb       = merge_roidb(roidbs)
    roidb       = filter_roidb(roidb, config)
    return roidb

def load_merge_filter_proposal_roidb_wrapper(config, proposal, output_path):
    """
    wrapper for load_proposal_roidb, merge_roidb, filter_roidb using config
    """
    image_sets = [iset for iset in config.dataset.image_set.split('+')]
    roidbs      = [load_proposal_roidb(config.dataset.dataset, image_set, config.dataset.root_path, 
            config.dataset.dataset_path, proposal=proposal, append_gt=True, flip=config.TRAIN.flip, result_path=output_path)
              for image_set in image_sets]
    roidb = merge_roidb(roidbs)
    roidb = filter_roidb(roidb, config)

def load_gt_roidb(dataset_name, image_set_name, root_path, dataset_path, result_path=None,
                  flip=False):
    """ load ground truth roidb """
    imdb = eval(dataset_name)(image_set_name, root_path, dataset_path, result_path)
    roidb = imdb.gt_roidb()
    if flip:
        roidb = imdb.append_flipped_images(roidb)
    return roidb


def load_proposal_roidb(dataset_name, image_set_name, root_path, dataset_path, result_path=None,
                        proposal='rpn', append_gt=True, flip=False):
    """ load proposal roidb (append_gt when training) """
    imdb = eval(dataset_name)(image_set_name, root_path, dataset_path, result_path)

    gt_roidb = imdb.gt_roidb()
    roidb = eval('imdb.' + proposal + '_roidb')(gt_roidb, append_gt)
    if flip:
        roidb = imdb.append_flipped_images(roidb)
    return roidb


def merge_roidb(roidbs):
    """ roidb are list, concat them together """
    roidb = roidbs[0]
    for r in roidbs[1:]:
        roidb.extend(r)
    return roidb


def filter_roidb(roidb, config):
    """ remove roidb entries without usable rois """

    def is_valid(entry):
        """ valid images have at least 1 fg or bg roi """
        overlaps = entry['max_overlaps']
        fg_inds = np.where(overlaps >= config.TRAIN.FG_THRESH)[0]
        bg_inds = np.where((overlaps < config.TRAIN.BG_THRESH_HI) & (overlaps >= config.TRAIN.BG_THRESH_LO))[0]
        valid = len(fg_inds) > 0 or len(bg_inds) > 0
        return valid

    num = len(roidb)
    filtered_roidb = [entry for entry in roidb if is_valid(entry)]
    num_after = len(filtered_roidb)
    print 'filtered %d roidb entries: %d -> %d' % (num - num_after, num, num_after)

    return filtered_roidb


def load_gt_segdb(dataset_name, image_set_name, root_path, dataset_path, result_path=None,
                  flip=False):
    """ load ground truth segdb """
    imdb = eval(dataset_name)(image_set_name, root_path, dataset_path, result_path)
    segdb = imdb.gt_segdb()
    if flip:
        segdb = imdb.append_flipped_images_for_segmentation(segdb)
    return segdb


def merge_segdb(segdbs):
    """ segdb are list, concat them together """
    segdb = segdbs[0]
    for r in segdbs[1:]:
        segdb.extend(r)
    return segdb
